import os
from tkinter import Tk, filedialog, messagebox, simpledialog, Label, Entry, Button, StringVar, ttk
from PIL import Image
import threading

def resize_images(source_folder, target_folder, max_size_kb, progress_var, success_var, fail_var, total_var):
    """
    压缩源文件夹及其子文件夹中的图片并保存到目标文件夹，保留文件的层级结构。

    :param source_folder: 源文件夹路径
    :param target_folder: 目标文件夹路径
    :param max_size_kb: 图片的最大文件大小，单位为KB
    :param progress_var: 进度条变量
    :param success_var: 成功数变量
    :param fail_var: 失败数变量
    :param total_var: 总数变量
    """
    # 确保目标文件夹存在
    if not os.path.exists(target_folder):
        os.makedirs(target_folder)

    # 遍历源文件夹及其子文件夹中的所有文件
    total_files = 0
    success_count = 0
    fail_count = 0
    for root, dirs, files in os.walk(source_folder):
        for filename in files:
            file_path = os.path.join(root, filename)

            # 检查文件是否为图片
            if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp')):
                total_files += 1

    # 更新总数
    total_var.set(total_files)

    # 创建线程列表
    threads = []

    # 遍历源文件夹及其子文件夹中的所有文件
    for index, (root, dirs, files) in enumerate(os.walk(source_folder)):
        for filename in files:
            file_path = os.path.join(root, filename)

            # 检查文件是否为图片
            if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                # 创建线程
                thread = threading.Thread(target=resize_image, args=(source_folder, file_path, target_folder, max_size_kb, progress_var, success_var, fail_var, total_var, index, total_files))
                threads.append(thread)
                thread.start()

    # 等待所有线程完成
    for thread in threads:
        thread.join()

    # 更新进度条为100%
    progress_var.set(100)

    # 显示完成消息
    messagebox.showinfo("提示", "图片压缩完成！")

def resize_image(source_folder, file_path, target_folder, max_size_kb, progress_var, success_var, fail_var, total_var, index, total_files):
    try:
        # 打开图片
        with Image.open(file_path) as img:
            # 计算相对路径
            relative_path = os.path.relpath(file_path, start=source_folder)
            # 构建目标文件路径
            target_file_path = os.path.join(target_folder, relative_path)
            # 确保目标目录存在
            os.makedirs(os.path.dirname(target_file_path), exist_ok=True)

            # 检查目标文件是否存在
            if os.path.exists(target_file_path):
                # 询问用户是否替换
                replace = messagebox.askyesno("提示", f"文件 {target_file_path} 已存在，是否替换？")
                if not replace:
                    return

            # 保存原始图片到临时文件
            temp_file_path = target_file_path
            img.save(temp_file_path)

            # 检查临时文件大小
            file_size_kb = os.path.getsize(temp_file_path) / 1024
            while file_size_kb > max_size_kb:
                # 计算缩放比例
                width, height = img.size
                ratio = 0.9  # 每次缩小10%
                new_size = (int(width * ratio), int(height * ratio))
                img = img.resize(new_size, Image.LANCZOS)

                # 保存压缩后的图片到临时文件
                img.save(temp_file_path)

                # 重新检查临时文件大小
                file_size_kb = os.path.getsize(temp_file_path) / 1024

            # 保存最终压缩后的图片到目标文件夹
            os.rename(temp_file_path, target_file_path)
            print(f"压缩并保存图片: {target_file_path}")
            # 更新成功数
            # 检查 success_var 是否为空，如果为空则初始化为 0
            if not success_var.get():
                success_var.set(0)
            success_count = int(success_var.get()) + 1
            success_var.set(success_count)
            # 更新进度条
            progress_var.set(success_count/ total_files * 100)
    except Exception as e:
        print(f"处理图片 {os.path.basename(file_path)} 时出错: {e}")
        # 检查 fail_var 是否为空，如果为空则初始化为 0
        if not fail_var.get():
            fail_var.set(0)
        fail_count = int(fail_var.get()) + 1
        fail_var.set(fail_count)
        # 更新进度条
        progress_var.set((success_count + fail_count) / total_files * 100)

    

def select_folder(entry_var, title):
    folder_path = filedialog.askdirectory(title=title)
    entry_var.set(folder_path)

def get_max_size(entry_var):
    max_size_kb = simpledialog.askinteger("输入", "请输入图片的最大文件大小（KB）：", minvalue=1)
    entry_var.set(max_size_kb)

def start_compression(source_folder, target_folder, max_size_kb):
    if not source_folder or not target_folder or not max_size_kb:
        messagebox.showinfo("提示", "请填写所有必填项。")
        return

    try:
        max_size_kb = int(max_size_kb)
        # 创建进度条和成功/失败数标签
        progress_var = StringVar()
        success_var = StringVar()
        fail_var = StringVar()
        total_var = StringVar()

        progress_bar = ttk.Progressbar(root, orient="horizontal", length=300, mode="determinate", variable=progress_var)
        progress_bar.grid(row=4, column=0, columnspan=3, pady=10)

        # 成功/失败数标签 默认值 0
        success_var.set(0)
        success_label = Label(root, text="成功数:")
        success_label.grid(row=5, column=0, sticky="e")
        success_count_label = Label(root, textvariable=success_var)
        success_count_label.grid(row=5, column=1, sticky="w")

        # 成功/失败数标签 默认值 0
        fail_var.set(0)
        fail_label = Label(root, text="失败数:")
        fail_label.grid(row=6, column=0, sticky="e")
        fail_count_label = Label(root, textvariable=fail_var)
        fail_count_label.grid(row=6, column=1, sticky="w")

        total_label = Label(root, text="总数:")
        total_label.grid(row=7, column=0, sticky="e")
        total_count_label = Label(root, textvariable=total_var)
        total_count_label.grid(row=7, column=1, sticky="w")

        # 启动图片压缩线程
        thread = threading.Thread(target=resize_images, args=(source_folder, target_folder, max_size_kb, progress_var, success_var, fail_var, total_var))
        thread.start()

    except ValueError:
        messagebox.showinfo("提示", "请输入有效的最大文件大小。")

def main():
    global root
    root = Tk()
    root.title("扬仔 x Waiting 图片压缩工具")

    # 源文件夹路径
    source_folder_var = StringVar()
    source_folder_label = Label(root, text="源文件夹路径:")
    source_folder_label.grid(row=0, column=0, sticky="e")
    source_folder_entry = Entry(root, textvariable=source_folder_var, width=50)
    source_folder_entry.grid(row=0, column=1, padx=10, pady=5)
    source_folder_button = Button(root, text="选择", command=lambda: select_folder(source_folder_var, "选择源文件夹"))
    source_folder_button.grid(row=0, column=2, padx=5, pady=5)

    # 目标文件夹路径
    target_folder_var = StringVar()
    target_folder_label = Label(root, text="目标文件夹路径:")
    target_folder_label.grid(row=1, column=0, sticky="e")
    target_folder_entry = Entry(root, textvariable=target_folder_var, width=50)
    target_folder_entry.grid(row=1, column=1, padx=10, pady=5)
    target_folder_button = Button(root, text="选择", command=lambda: select_folder(target_folder_var, "选择目标文件夹"))
    target_folder_button.grid(row=1, column=2, padx=5, pady=5)

    # 最大文件大小
    max_size_kb_var = StringVar()
    max_size_kb_label = Label(root, text="最大文件大小（KB）:")
    max_size_kb_label.grid(row=2, column=0, sticky="e")
    max_size_kb_entry = Entry(root, textvariable=max_size_kb_var, width=10)
    max_size_kb_entry.grid(row=2, column=1, padx=10, pady=5)
    max_size_kb_button = Button(root, text="设置", command=lambda: get_max_size(max_size_kb_var))
    max_size_kb_button.grid(row=2, column=2, padx=5, pady=5)

    # 开始压缩按钮
    start_button = Button(root, text="开始压缩", command=lambda: start_compression(source_folder_var.get(), target_folder_var.get(), max_size_kb_var.get()))
    start_button.grid(row=3, column=1, pady=10)

    root.mainloop()

if __name__ == "__main__":
    main()
